# -*- coding: UTF-8 -*-
# *******************************************************************
# File Name: background_img_text
# > Author: 04000387
# > Created Time: 2025/1/6 14:34
# *******************************************************************
from PIL import Image, ImageDraw, ImageFont
import cv2
import numpy as np
from pathlib import Path
import json


def _calculate_luminance(_color):
    """计算颜色的亮度"""
    return 0.299 * _color[0] + 0.587 * _color[1] + 0.114 * _color[2]


def _get_contrast_color(background_color):
    """根据背景色选择对比度较高的字体颜色"""
    luminance = _calculate_luminance(background_color)
    if luminance > 128:
        # 背景较亮，选择较暗的字体颜色
        return np.random.randint(0, 127), np.random.randint(0, 127), np.random.randint(0, 127)
    else:
        # 背景较暗，选择较亮的字体颜色
        return np.random.randint(128, 255), np.random.randint(128, 255), np.random.randint(128, 255)


def _generate_img_text(_img_path, _text, _font_size=25, _out_put_dir="../generate/img/img_text"):
    _out_put_dir = Path(_out_put_dir)
    if not _out_put_dir.exists():
        _out_put_dir.mkdir(parents=True)

    bk_img = cv2.imread(_img_path)
    # 设置需要显示的字体
    font_path = "font/simsun.ttc"
    font = ImageFont.truetype(font_path, _font_size)
    img_pil = Image.fromarray(bk_img)
    draw = ImageDraw.Draw(img_pil)

    h, w, _ = bk_img.shape
    used_text = []

    line_num = h // (_font_size + 3)
    col_num = w // (_font_size + 3)

    split_texts = [it for it in _text]
    # 绘制文字信息
    init_height = 0
    while True:
        dh = np.random.randint(init_height, line_num)

        init_height += dh
        h_pix = init_height * (_font_size + 3)
        if (h_pix + _font_size) >= h or init_height >= line_num:
            break

        dw = np.random.randint(col_num - 1)

        char_num = col_num - dw - 1
        one_text = np.random.choice(split_texts, char_num, replace=True)
        used_text.append("".join(one_text))
        w_pix = dw * (_font_size + 3)

        (b, g, r) = bk_img[h_pix, w_pix]

        font_color = _get_contrast_color((r, g, b))
        draw.text((w_pix, h_pix), used_text[-1], font=font, fill=font_color, stroke_width=1)
        bk_img = np.array(img_pil)

    cv2.imshow("add_text", bk_img)
    cv2.waitKey()
    _save_path = _out_put_dir / _img_path.name
    cv2.imwrite(str(_save_path.absolute()), bk_img)

    return "\n".join(used_text)


def _get_files(_root_dir, suffix=None, _num=-1):
    if suffix is None:
        suffix = [".png", ".jpg"]
    if isinstance(suffix, str):
        suffix = [suffix]

    _root_dir = Path(_root_dir)
    _root_dir.is_dir()
    if _root_dir.is_file():
        return [_root_dir]

    _dirs = [_root_dir]

    _res = []

    while True:
        if len(_dirs) == 0:
            break

        _dir = _dirs.pop(0)
        for it in _dir.iterdir():
            if it.is_file() and it.suffix in suffix:
                _res.append(it)

            elif it.is_dir():
                _dirs.append(it)

            if 0 < _num <= len(_res):
                break

    return _res


def _get_texts(_text_folder):
    _text_files = _get_files(_text_folder, ".txt")
    _res = []

    for _it in _text_files:
        _fp = _it.open("r", encoding="utf-8")
        for _line in _fp.readlines():
            if len(_line.strip().strip("\n")) <= 1:
                continue
            _res.append(_line)

        _fp.close()

    return _res


def generate_img_text(image_folder, text_folder, out_path="../generate/generate_text_img", font_size=25, _img_num=2):
    files = _get_files(image_folder, _num=_img_num)
    texts = _get_texts(text_folder)
    out_path = Path(out_path)

    if not out_path.exists():
        out_path.mkdir(parents=True)
        (out_path / "img").mkdir(parents=True)

    text_jsonl = out_path / "img_text.jsonl"
    img_path = out_path / "img"

    fp = text_jsonl.open("w", encoding="utf-8")
    texts_write = []
    for _img in files:
        text = _generate_img_text(_img, np.random.choice(texts, replace=False), font_size,
                                  _out_put_dir=str(img_path.absolute()))

        string = "这是一张带有文字的图片,图片中的文字如下:" + text
        one = json.dumps({"img_file": str(img_path), "content": string}) + "\n"
        texts_write.append(one)

    fp.writelines(texts_write)

    fp.close()
